import numpy as np  #提供了处理数组和数值计算的功能
import pandas as pd#用于数据处理和分析，提供了高性能、易用的数据结构和数据分析工具
import torch#PyTorch 深度学习框架的主要库
import torch.nn as nn # 构建和训练神经网络
import torch.optim as optim # 优化器，用于调整模型的权重和学习率等超参数
from sklearn.model_selection import train_test_split#划分数据集成训练集和测试集
from sklearn.preprocessing import StandardScaler # 用于数据标准化，将特征数据按照标准正态分布进行缩放
import matplotlib.pyplot as plt#绘制数据可视化图形
from sklearn.metrics import r2_score

# 固定随机种子
seed = 10  #设置了随机种子，以确保每次运行程序得到的随机数都是一样的
np.random.seed(seed)
torch.manual_seed(seed)


data = pd.read_csv('input.csv', encoding="gbk")#使用 pandas 读取数据集，并将其中的特征和标签分别赋值给变量 X 和 y
cols = ['刀盘转速', '刀盘转速设定值', '推进速度', '推进速度设定值', '刀盘推力', '刀盘扭矩']
data = data[cols]
data = data.dropna(how='any') #删除缺失数据的行
X = data[['刀盘转速', '刀盘转速设定值', '推进速度', '推进速度设定值',  '刀盘扭矩']]
y = data[['刀盘推力']]

# 特征标准化
scaler = StandardScaler()
X = scaler.fit_transform(X)

# 转换数据为PyTorch张量
X = torch.tensor(X, dtype=torch.float32)
y = torch.tensor(y.values, dtype=torch.float32)

# 数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=seed)#train_test_split 函数将数据集划分为训练集和测试集，测试集大小为数据集大小的 20%，训练集大小为数据集大小的 80%。


# 创建前馈神经网络模型
class FNN(nn.Module):#它继承了 PyTorch 的 nn.Module 类
    def __init__(self, input_dim):
        super(FNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 64)#第一个线性层输入维度为 input_dim，输出维度为 64
        self.fc2 = nn.Linear(64, 64)#第二个线性层输入维度为 64，输出维度为 64
        self.fc3 = nn.Linear(64, 32)
        self.fc4 = nn.Linear(32, 1)#第三个线性层输入维度为 32，输出维度为 1

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = self.fc4(x)
        return x

# 模型实例化，并定义损失函数和优化器
model = FNN(X_train.shape[1])
criterion = nn.MSELoss() # 均方根损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
epochs = 800 # 模型迭代次数
train_losses = []  # 创建两个空列表，用于存储每个迭代周期的训练损失和测试损失
test_losses = []
r_squareds = []
for epoch in range(epochs):
    model.train()  # 训练模式
    optimizer.zero_grad()  # 清空优化器的梯度
    outputs = model(X_train)

    loss = criterion(outputs, y_train)  # 预测输出和真实标签 y_train 计算损失值 los
    loss.backward()  # 使用 loss.backward()，对损失值进行反向传播
    optimizer.step()
    train_losses.append(loss.item())  # 将每个迭代周期的训练损失 loss.item() 添加到 train_losses 列表中

    model.eval()
    with torch.no_grad():  # 将测试数据 X_test 输入模型，得到模型的预测输出 test_outputs
        test_outputs = model(X_test)
        test_loss = criterion(test_outputs, y_test)  # 根据预测输出和真实标签 y_test 计算测试损失值 test_loss
        test_losses.append(test_loss.item())  # 将每个迭代周期的测试损失 test_loss.item() 添加到 test_losses 列表中
        r_squared = r2_score(y_test,test_outputs.detach().numpy())
        r_squareds.append(r_squared)

    print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item()}, r_squared: {r_squared} ')  # 打印当前迭代周期的序号和损失值


# 绘制训练集和测试集的损失曲线
#通过绘制曲线将训练集和测试集的损失随迭代次数的变化情况可视化出来，以便观察模型的训练过程和性能

plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
plt.figure(figsize=(10, 6))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs + 1), train_losses, label='训练集损失', marker='o')
plt.plot(range(1, epochs + 1), test_losses, label='测试集损失', marker='s')
plt.xlabel("Epochs")
plt.ylabel("Loss (MSE)")
plt.title("训练集和测试集损失曲线")
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(range(350, epochs + 1),r_squareds[349:], label='R-Square')
plt.legend()
plt.title('R-Square曲线')
plt.xlabel('训练轮数')
plt.ylabel('R-Square值')

plt.tight_layout()
plt.show()

